// This test reproduces the test problem 3 from
// 'A Cartesian Grid Embedded Boundary Method for Poisson's Equation on Irregular Domains'
// (Johansen and Colella 1998),
// but with fourth-order accuracy.

#include <iostream>
#include <iomanip>
#include <vector>
#include <array>
#include "Core/numlib.h"
#include "Core/TensorExpr.h"
#include "FiniteDiff/LevelOp.h"
#include "FiniteDiff/Multigrid.h"
#include "Embedded/EmbeddedGrid.h"
#include "Embedded/EmLevelOp.h"
#include "Embedded/EmFuncFiller.h"
#include "Embedded/EmIntergrid.h"
#include "Embedded/PoissonOp.h"
#include "Embedded/BlockRelaxation.h"
#include "../example/common/Ellipsoid.h"
#include "../example/common/RoundedStar.h"
#include "../example/common/TestUtility.h"

#if DIM == 2

const int Dim = 2;
using rVec = Vec<Real, Dim>;
template class Tensor<Real, 2>;


//======================================================================

const char *bcTypes = "DDDDN";

class PoissonLevelOp : public MGLevelOp<2>
{
public:
  enum { Dim = 2 };
  using EmGrid = EmbeddedGrid<Dim>;
  using BaseClass = MGLevelOp<Dim>;

  PoissonLevelOp(const EmGrid &aGrid, Real w4j, const char *aBcTypes, bool isBottom = false)
      : BaseClass(aGrid, w4j, aBcTypes), lvOp(aGrid), eff(aGrid), elv(aGrid),
        poissonOp(aGrid, bcTypes)
  {
    pbr = new BlockRelaxation<Dim>(aGrid, poissonOp.matOp, 0.5);
    if(isBottom)
      pBottomSolver = new PoissonBottomSolver<Dim>(poissonOp);
    else
      pBottomSolver = nullptr;
  }

  virtual ~PoissonLevelOp() {
    delete pBottomSolver;
    delete pbr;
  }

  Real computeNorm(const Tensor<Real, Dim> &phi) const override
  {
    return elv.NRM(phi, CellCentered, 0);
  }

  void fillGhosts(Tensor<Real, Dim> &phi, const Tensor<Real, Dim> &rhs, const Tensor<Real, Dim-1> *bcData) const override
  {
    BaseClass::fillGhosts(phi, rhs, bcData);
    eff.transferBdryValues(phi, bcData[2*Dim]);
    pbr->apply(phi, rhs);
  }

  void fillGhosts(Tensor<Real, Dim> &phi, const Tensor<Real, Dim> &rhs, int) const override
  {
    BaseClass::fillGhosts(phi, rhs, 0);
    eff.transferBdryValues(phi, 0);
    pbr->apply(phi, rhs);
  }

  void applySmoother(const Tensor<Real, Dim> &phi, Tensor<Real, Dim> &smoothed, const Tensor<Real, Dim> &rhs) const override
  {
    lvOp.RELAX_4TH(phi, rhs, smoothed, weightForJacobi);
  }

  void computeResidual(const Tensor<Real, Dim> &phi, Tensor<Real, Dim> &rsd, const Tensor<Real, Dim> &rhs) const override
  {
    poissonOp.apply(phi, rhs, rsd);
  }

  void solveBottom(Tensor<Real, Dim> &phi, const Tensor<Real, Dim> &rhs, int numBottomRelax) const override
  {
    assert(pBottomSolver != nullptr);
    pBottomSolver->solve(phi, rhs, 0);
  }

  const PoissonOp<Dim> &getPoissonOp() const { return poissonOp; }

  using MGLevelOp<2>::getDomain;

protected:
  LevelOp<Dim> lvOp;
  EmFuncFiller<Dim> eff;
  EmLevelOp<Dim> elv;
  PoissonOp<Dim> poissonOp;
  PoissonBottomSolver<Dim> *pBottomSolver;
  BlockRelaxation<Dim> *pbr;
  using MGLevelOp<Dim>::weightForJacobi;
};

const PoissonLevelOp *pDbgLevelOp;

template <class T0, class T1, class T2>
std::array<Real,3> testNeumannPoisson(const std::vector<const EmbeddedGrid<2> *> &gridHier,
                                      const std::vector<const MGLevelOp<2> *> &opHier,
                                      const std::vector<const Intergrid<2> *> &trHier,
                                      const T0 &PHI, const T1 &DPHI, const T2 &D2PHI)
{
  const EmbeddedGrid<Dim> &emGrid = *gridHier.front();
  EmLevelOp<Dim> elv(emGrid);
  EmFuncFiller<Dim> eff(emGrid);
  const Box<Dim> gbx = emGrid.getGhostedBox();
  Tensor<Real, Dim> phi(gbx);
  Tensor<Real, Dim> err(gbx);
  Tensor<Real, Dim> rhs(gbx);
  Tensor<Real, Dim-1> cData[2*Dim+1];
  eff.fillPointValues(err, CellCentered, PHI);
  eff.fillPointValues(rhs, CellCentered, D2PHI);
  phi = 0.0;
  for(int d = 0; d < Dim; ++d)
    for (int m : {0, 1})
      cData[2 * d + m].resize(reduce(gbx, d));
  eff.fillBdryValues(cData, CellCentered, PHI);
  cData[2*Dim].resize(emGrid.bdryNodes.size());
//  eff.fillBdryValues(cData[2*Dim], "curvilinear", PHI); // Dirichlet boundary
  eff.fillBdryGrads(cData[2*Dim], "curvilinear", DPHI); // Neumann boundary

  if(true) { // test the solution error
    Multigrid<Dim> mg(opHier, trHier);
    mg.setParam({3, 3, 0}, 40, 0.0, 1e-13);
    CPUTimer timer;
    timer.reset();
    mg.solve(phi, rhs, cData);
    std::cout << "Multigrid : " << std::defaultfloat << std::setprecision(5)
              << timer() << "s. " << std::endl;
    elv.zeroExtNodes(phi);
    err = err - phi;
  } else { // test the truncation error
    eff.transferBdryValues(err, cData[2*Dim]);
    const PoissonOp<Dim> &poissonOp = dynamic_cast<const PoissonLevelOp *>(opHier.front())->getPoissonOp();
    poissonOp.apply(err, phi);
    err = rhs - phi;
  }
  elv.zeroExtNodes(err);

  std::array<Real, 3> errnorm {};
  for(int p = 0; p <= 2; ++p)
    errnorm[p] = elv.NRM(err, CellCentered, p);
  return errnorm;
}

template <class T0, class T1, class T2>
void doTest(const std::vector<int> &gridSize, int numCompHier, const T0 &PHI, const T1 &DPHI, const T2 &D2PHI)
{
  const Real tol = 1e-12;
  const int numGrid = gridSize.size();
  std::vector<const EmbeddedGrid<Dim> *> gridHier;
  std::vector<const PoissonLevelOp *> opHier;
  std::vector<const EmIntergrid<Dim> *> trHier;
  rVec translate { 0.5, 0.5 };
  RoundedStar geoObj(0.25, 0.05, 6);
//  Ellipsoid<Dim> geoObj(0.5, {0.25, 0.25});
  for(int n = 0; n < numGrid; ++n) {
    const int N = gridSize[n];
    std::cout << "\n\nInitializing h = " << oneOver(N) << std::endl;
    RectDomain<Dim> rd(Box<Dim>(0, N-1), 1.0/N, CellCentered, 2);
    auto *pGrid = new EmbeddedGrid<Dim>(
        rd,
        [&](rVec q, rVec &closest, rVec &normal, bool &contained, Real tol) {
          q = q - translate;
          geoObj.check(q, closest, normal, contained, tol);
          closest = closest + translate;
          normal = -normal;
          contained = !contained;
        }, true, 0.1, tol);
#ifndef NDEBUG
    if(get_dbglevel() >= 2)
      pGrid->dump(tmpos);
#endif
    gridHier.push_back(pGrid);
    push_dbglevel(1);
    auto *pPoissonOp = new PoissonLevelOp(*pGrid, 0.6, bcTypes, n == numGrid-1);
    std::cout.flush();
    pop_dbglevel();
    opHier.push_back(pPoissonOp);
    if(n != 0) {
      auto *pT = new EmIntergrid<Dim>(*pGrid, *(gridHier[n-1]));
      trHier.push_back(pT);
    }
  }

  std::vector<std::array<Real, 3>> errnorm(numCompHier);
  for(int n = numCompHier-1; n >= 0; --n) {
    std::vector<const EmbeddedGrid<Dim> *> _gridHier(gridHier.cbegin() + n, gridHier.cend());
    std::vector<const MGLevelOp<Dim> *> _opHier(opHier.cbegin() + n, opHier.cend());
    std::vector<const Intergrid<Dim> *> _trHier(trHier.cbegin() + n, trHier.cend());
    std::cout << "\n\nTesting h = " << oneOver(gridSize[n]) << std::endl;
    errnorm[n] = testNeumannPoisson(_gridHier, _opHier, _trHier, PHI, DPHI, D2PHI);
  }

  for(int n = numGrid-1; n >= 0; --n) {
    if(n != 0)
      delete trHier[n-1];
    delete opHier[n];
    delete gridHier[n];
  }
  printConvergenceTable(&gridSize[0], errnorm);
}

int main()
{
  reset_dbglevel(1);
  std::cout << "SpaceDim = " << SpaceDim << std::endl;
  std::cout << "============================================================================\n";
  // the analytic solutions in polar coordinates
  auto POLAR_PHI = [](Real r, Real theta) { return ipow<4>(r) * cos(3*theta); };
  auto POLAR_DPHI = [](Real r, Real theta) {
    Real dr = 4 * ipow<3>(r) * cos(3*theta);
    Real dt = -3.0 * ipow<4>(r) * sin(3*theta);
    // return the gradient in (x,y) coordinate
    return Vec<Real, 2> { cos(theta) * dr - sin(theta)*dt/r,
                          sin(theta) * dr + cos(theta)*dt/r };
  };
  auto POLAR_D2PHI = [](Real r, Real theta) { return 7.0 * ipow<2>(r) * cos(3*theta); };
  // the analytic solutions in Cartesian coordinates
  Vec<Real, 2> translate {0.5, 0.5};
  auto PHI = [&](rVec x) {
//    return sin(2*M_PI*x[0]) * cos(2*M_PI*x[1]);
    x = x - translate;
    return POLAR_PHI(norm(x, 2), atan2(x[1], x[0]));
  };
  auto DPHI = [&](rVec x) {
//    return rVec { cos(2*M_PI*x[0]) * cos(2*M_PI*x[1]),
//                  -sin(2*M_PI*x[0]) * sin(2*M_PI*x[1]) } * (2*M_PI);
    x = x - translate;
    return POLAR_DPHI(norm(x, 2), atan2(x[1], x[0]));
  };
  auto D2PHI = [&](rVec x) {
//    return -8*M_PI*M_PI * sin(2*M_PI*x[0]) * cos(2*M_PI*x[1]);
    x = x - translate;
    return POLAR_D2PHI(norm(x, 2), atan2(x[1], x[0]));
  };
  //
  doTest({320, 160, 80, 40}, 4, PHI, DPHI, D2PHI);
  std::cout << std::endl;
  return 0;
}

#else

int main()
{
  std::cout << "============================================================================\n";
  std::cout << "testNeumannPoisson is skipped for 3D. " << std::endl;
  return 0;
}

#endif // DIM == 2
